import torch
import torch.nn as nn
import torch.optim as optim
import random
import numpy as np
import time
import pygame
import sys
from collections import deque

# 初始化Pygame
pygame.init()

# 棋盘常量定义
BOARD_ROWS = 10
BOARD_COLS = 9
CELL_SIZE = 60  # 每个格子的像素大小
BOARD_WIDTH = CELL_SIZE * BOARD_COLS
BOARD_HEIGHT = CELL_SIZE * BOARD_ROWS
SCREEN_WIDTH = BOARD_WIDTH
SCREEN_HEIGHT = BOARD_HEIGHT + 50  # 底部留一些空间显示信息

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (200, 0, 0)
BLUE = (0, 0, 200)
YELLOW = (255, 255, 0)
BROWN = (139, 69, 19)
LIGHT_BROWN = (222, 184, 135)

# 创建屏幕
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("中国象棋 - 人机对战")

# 加载字体
try:
    font = pygame.font.Font("simhei.ttf", 24)  # 尝试加载黑体
except:
    font = pygame.font.SysFont(["SimHei", "WenQuanYi Micro Hei", "Heiti TC"], 24)

# 棋子编码（正为红方，负为黑方）
PIECE_CODES = {
    '红帅': 1, '红仕': 2, '红相': 3, '红马': 4, '红车': 5, '红炮': 6, '红兵': 7,
    '黑将': -1, '黑士': -2, '黑象': -3, '黑马': -4, '黑车': -5, '黑炮': -6, '黑卒': -7,
    '空': 0
}
PIECE_NAMES = {v: k for k, v in PIECE_CODES.items()}
PIECE_CHARS = {  # 棋子显示的字符
    1: '帅', 2: '仕', 3: '相', 4: '马', 5: '车', 6: '炮', 7: '兵',
    -1: '将', -2: '士', -3: '象', -4: '马', -5: '车', -6: '炮', -7: '卒'
}

# 特殊位置定义
RED_PALACE = [(0, 3), (0, 4), (0, 5), (1, 4), (2, 3), (2, 4), (2, 5)]  # 红方九宫
BLACK_PALACE = [(7, 3), (7, 4), (7, 5), (8, 4), (9, 3), (9, 4), (9, 5)]  # 黑方九宫
RIVER_LINE = 5  # 过河线（红方过线y≥5，黑方过线y≤4）


class ChineseChessBoard:
    """棋盘核心类，处理棋盘状态和规则逻辑"""
    def __init__(self):
        self.board = torch.zeros((BOARD_ROWS, BOARD_COLS), dtype=torch.int8)  # 10x9棋盘
        self.current_player = 1  # 1=红方（玩家），-1=黑方（AI）
        self.history = []  # 记录历史棋局
        self.selected_pos = None  # 当前选中的棋子位置
        self.valid_moves = []  # 选中棋子的有效移动位置
        self.init_board()

    def init_board(self):
        """初始化标准棋盘布局"""
        # 红方布局（y=0：车马相仕帅仕相马车；y=2：兵）
        self.board[0] = torch.tensor([5, 4, 3, 2, 1, 2, 3, 4, 5])
        self.board[2] = torch.tensor([0, 7, 0, 7, 0, 7, 0, 7, 0])
        # 黑方布局（y=9：车马象士将士象马车；y=7：卒）
        self.board[9] = torch.tensor([-5, -4, -3, -2, -1, -2, -3, -4, -5])
        self.board[7] = torch.tensor([0, -7, 0, -7, 0, -7, 0, -7, 0])
        self.history.append(self.board.clone())

    def draw_board(self):
        """绘制棋盘"""
        # 绘制背景
        screen.fill(BROWN)
        
        # 绘制棋盘网格
        for y in range(BOARD_ROWS):
            for x in range(BOARD_COLS):
                # 绘制网格线
                pygame.draw.rect(screen, BLACK, 
                                (x*CELL_SIZE, y*CELL_SIZE, CELL_SIZE, CELL_SIZE), 1)
        
        # 绘制楚河汉界
        river_text = font.render("楚 河        汉 界", True, BLACK)
        screen.blit(river_text, (BOARD_WIDTH//2 - river_text.get_width()//2, 
                                CELL_SIZE*4 + CELL_SIZE//2 - river_text.get_height()//2))
        
        # 绘制九宫格斜线
        # 红方九宫
        pygame.draw.line(screen, BLACK, 
                        (3*CELL_SIZE, 0*CELL_SIZE), (5*CELL_SIZE, 2*CELL_SIZE), 2)
        pygame.draw.line(screen, BLACK, 
                        (5*CELL_SIZE, 0*CELL_SIZE), (3*CELL_SIZE, 2*CELL_SIZE), 2)
        # 黑方九宫
        pygame.draw.line(screen, BLACK, 
                        (3*CELL_SIZE, 7*CELL_SIZE), (5*CELL_SIZE, 9*CELL_SIZE), 2)
        pygame.draw.line(screen, BLACK, 
                        (5*CELL_SIZE, 7*CELL_SIZE), (3*CELL_SIZE, 9*CELL_SIZE), 2)

        # 绘制棋子
        for y in range(BOARD_ROWS):
            for x in range(BOARD_COLS):
                piece = self.board[y, x].item()
                if piece != 0:
                    # 棋子颜色（红方/黑方）
                    color = RED if piece > 0 else BLACK
                    bg_color = WHITE
                    
                    # 绘制棋子
                    pygame.draw.circle(screen, bg_color, 
                                     (x*CELL_SIZE + CELL_SIZE//2, y*CELL_SIZE + CELL_SIZE//2), 
                                     CELL_SIZE//2 - 5)
                    pygame.draw.circle(screen, color, 
                                     (x*CELL_SIZE + CELL_SIZE//2, y*CELL_SIZE + CELL_SIZE//2), 
                                     CELL_SIZE//2 - 5, 2)
                    
                    # 绘制棋子文字
                    text = font.render(PIECE_CHARS[piece], True, color)
                    text_rect = text.get_rect(center=(x*CELL_SIZE + CELL_SIZE//2, 
                                                     y*CELL_SIZE + CELL_SIZE//2))
                    screen.blit(text, text_rect)
        
        # 绘制选中效果
        if self.selected_pos:
            y, x = self.selected_pos
            # 选中棋子高亮
            pygame.draw.circle(screen, YELLOW, 
                             (x*CELL_SIZE + CELL_SIZE//2, y*CELL_SIZE + CELL_SIZE//2), 
                             CELL_SIZE//2 - 5, 3)
            
            # 绘制有效移动位置标记
            for (ny, nx) in self.valid_moves:
                pygame.draw.circle(screen, YELLOW, 
                                 (nx*CELL_SIZE + CELL_SIZE//2, ny*CELL_SIZE + CELL_SIZE//2), 
                                 10)

        # 绘制状态信息
        status_text = "你的回合（红方）" if self.current_player == 1 else "AI思考中..."
        text_surface = font.render(status_text, True, WHITE)
        screen.blit(text_surface, (10, BOARD_HEIGHT + 10))

        pygame.display.flip()

    def is_in_palace(self, y, x, player):
        """判断坐标是否在当前玩家的九宫内"""
        return (y, x) in RED_PALACE if player == 1 else (y, x) in BLACK_PALACE

    def get_piece_moves(self, y, x):
        """获取指定位置棋子的所有合法移动"""
        piece = self.board[y, x].item()
        if piece == 0 or piece * self.current_player <= 0:
            return []  # 空位置或对方棋子无移动权限

        moves = []
        abs_piece = abs(piece)

        # 1. 车：直线移动（水平/垂直无阻挡）
        if abs_piece == 5:
            # 水平方向
            for dx in [1, -1]:
                new_x = x + dx
                while 0 <= new_x < BOARD_COLS:
                    if self.board[y, new_x] == 0:
                        moves.append((y, new_x))
                    else:
                        if self.board[y, new_x] * self.current_player < 0:
                            moves.append((y, new_x))  # 可吃对方棋子
                        break
                    new_x += dx
            # 垂直方向
            for dy in [1, -1]:
                new_y = y + dy
                while 0 <= new_y < BOARD_ROWS:
                    if self.board[new_y, x] == 0:
                        moves.append((new_y, x))
                    else:
                        if self.board[new_y, x] * self.current_player < 0:
                            moves.append((new_y, x))  # 可吃对方棋子
                        break
                    new_y += dy

        # 2. 马：日字走法（防蹩马腿）
        elif abs_piece == 4:
            horse_moves = [
                (y+2, x+1), (y+2, x-1), (y-2, x+1), (y-2, x-1),
                (y+1, x+2), (y+1, x-2), (y-1, x+2), (y-1, x-2)
            ]
            block_pos = [  # 蹩马腿的位置
                (y+1, x), (y+1, x), (y-1, x), (y-1, x),
                (y, x+1), (y, x-1), (y, x+1), (y, x-1)
            ]
            for i, (ny, nx) in enumerate(horse_moves):
                if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS:
                    by, bx = block_pos[i]
                    if self.board[by, bx] == 0 and self.board[ny, nx] * self.current_player <= 0:
                        moves.append((ny, nx))

        # 3. 炮：直线移动（吃子需隔一子）
        elif abs_piece == 6:
            # 水平方向
            for dx in [1, -1]:
                new_x = x + dx
                has_block = False
                while 0 <= new_x < BOARD_COLS:
                    if self.board[y, new_x] == 0 and not has_block:
                        moves.append((y, new_x))
                    elif self.board[y, new_x] != 0 and not has_block:
                        has_block = True  # 记录第一个阻挡
                    elif self.board[y, new_x] != 0 and has_block:
                        if self.board[y, new_x] * self.current_player < 0:
                            moves.append((y, new_x))  # 隔子吃对方
                        break
                    new_x += dx
            # 垂直方向
            for dy in [1, -1]:
                new_y = y + dy
                has_block = False
                while 0 <= new_y < BOARD_ROWS:
                    if self.board[new_y, x] == 0 and not has_block:
                        moves.append((new_y, x))
                    elif self.board[new_y, x] != 0 and not has_block:
                        has_block = True
                    elif self.board[new_y, x] != 0 and has_block:
                        if self.board[new_y, x] * self.current_player < 0:
                            moves.append((new_y, x))  # 隔子吃对方
                        break
                    new_y += dy

        # 4. 相/象：田字走法（不过河，防塞象眼）
        elif abs_piece == 3:
            if (self.current_player == 1 and y >= RIVER_LINE) or (self.current_player == -1 and y < RIVER_LINE):
                return []  # 红相不过河，黑象不过河
            elephant_moves = [(y+2, x+2), (y+2, x-2), (y-2, x+2), (y-2, x-2)]
            block_pos = [(y+1, x+1), (y+1, x-1), (y-1, x+1), (y-1, x-1)]
            for i, (ny, nx) in enumerate(elephant_moves):
                if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS:
                    by, bx = block_pos[i]
                    if self.board[by, bx] == 0 and self.board[ny, nx] * self.current_player <= 0:
                        moves.append((ny, nx))

        # 5. 仕/士：九宫内斜线走法
        elif abs_piece == 2:
            scholar_moves = [(y+1, x+1), (y+1, x-1), (y-1, x+1), (y-1, x-1)]
            for ny, nx in scholar_moves:
                if self.is_in_palace(ny, nx, self.current_player) and self.board[ny, nx] * self.current_player <= 0:
                    moves.append((ny, nx))

        # 6. 帅/将：九宫内直线走法（不照面）
        elif abs_piece == 1:
            king_moves = [(y+1, x), (y-1, x), (y, x+1), (y, x-1)]
            for ny, nx in king_moves:
                if self.is_in_palace(ny, nx, self.current_player) and self.board[ny, nx] * self.current_player <= 0:
                    if not self.is_king_face_to_face(ny, nx):
                        moves.append((ny, nx))

        # 7. 兵/卒：未过河前直走，过河后可左右
        elif abs_piece == 7:
            if self.current_player == 1:  # 红兵
                if y < RIVER_LINE:  # 未过河（只能向前）
                    if y+1 < BOARD_ROWS and self.board[y+1, x] * self.current_player <= 0:
                        moves.append((y+1, x))
                else:  # 过河后（向前+左右）
                    for (dy, dx) in [(1, 0), (0, 1), (0, -1)]:
                        ny, nx = y+dy, x+dx
                        if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS and self.board[ny, nx] * self.current_player <= 0:
                            moves.append((ny, nx))
            else:  # 黑卒
                if y >= RIVER_LINE - 1:  # 未过河（只能向前）
                    if y-1 >= 0 and self.board[y-1, x] * self.current_player <= 0:
                        moves.append((y-1, x))
                else:  # 过河后（向前+左右）
                    for (dy, dx) in [(-1, 0), (0, 1), (0, -1)]:
                        ny, nx = y+dy, x+dx
                        if 0 <= ny < BOARD_ROWS and 0 <= nx < BOARD_COLS and self.board[ny, nx] * self.current_player <= 0:
                            moves.append((ny, nx))

        return moves

    def is_king_face_to_face(self, king_y, king_x):
        """判断双方帅/将是否照面（同一列无阻挡）"""
        enemy_king = -1 if self.current_player == 1 else 1
        enemy_pos = torch.where(self.board == enemy_king)
        if len(enemy_pos[0]) == 0:
            return False  # 对方帅/将已被吃
        ey, ex = enemy_pos[0].item(), enemy_pos[1].item()
        if king_x != ex:
            return False  # 不同列不照面
        # 检查中间是否有阻挡
        min_y, max_y = min(king_y, ey), max(king_y, ey)
        for y in range(min_y + 1, max_y):
            if self.board[y, king_x].item() != 0:
                return False
        return True

    def is_checkmate(self):
        """判断当前玩家是否被将死（无合法移动）"""
        for y in range(BOARD_ROWS):
            for x in range(BOARD_COLS):
                if self.board[y, x].item() * self.current_player > 0:
                    if len(self.get_piece_moves(y, x)) > 0:
                        return False  # 还有合法移动
        return True

    def move_piece(self, from_y, from_x, to_y, to_x):
        """执行落子（合法返回True）"""
        # 边界校验
        if not (0 <= from_y < BOARD_ROWS and 0 <= from_x < BOARD_COLS):
            return False
        if not (0 <= to_y < BOARD_ROWS and 0 <= to_x < BOARD_COLS):
            return False
        # 棋子所有权校验
        if self.board[from_y, from_x].item() * self.current_player <= 0:
            return False
        # 移动合法性校验
        if (to_y, to_x) not in self.get_piece_moves(from_y, from_x):
            return False

        # 执行落子
        self.history.append(self.board.clone())
        self.board[to_y, to_x] = self.board[from_y, from_x]
        self.board[from_y, from_x] = 0
        self.current_player *= -1  # 切换玩家
        self.selected_pos = None  # 清除选中状态
        self.valid_moves = []
        return True

    def handle_click(self, pos):
        """处理鼠标点击事件"""
        if self.current_player != 1:  # 只有玩家回合响应点击
            return False
            
        x, y = pos
        # 转换为棋盘坐标
        board_x = x // CELL_SIZE
        board_y = y // CELL_SIZE
        
        # 检查是否在棋盘范围内
        if not (0 <= board_x < BOARD_COLS and 0 <= board_y < BOARD_ROWS):
            return False
            
        # 如果点击了自己的棋子，选中它并显示可移动位置
        piece = self.board[board_y, board_x].item()
        if piece * self.current_player > 0:
            self.selected_pos = (board_y, board_x)
            self.valid_moves = self.get_piece_moves(board_y, board_x)
            return True
            
        # 如果已经选中棋子，尝试移动
        elif self.selected_pos:
            from_y, from_x = self.selected_pos
            if (board_y, board_x) in self.valid_moves:
                # 执行移动
                result = self.move_piece(from_y, from_x, board_y, board_x)
                return result
                
        # 取消选中
        self.selected_pos = None
        self.valid_moves = []
        return False


class ChessValueNet(nn.Module):
    """价值网络：预测当前棋局的胜率（-1~1）"""
    def __init__(self):
        super().__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.ReLU()
        )
        self.fc_layers = nn.Sequential(
            nn.Flatten(),
            nn.Linear(128 * 2 * 2, 256),  # 卷积后尺寸：10→5→2，9→4→2
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, 1),
            nn.Tanh()  # 输出范围[-1,1]（1=红方胜，-1=黑方胜）
        )

    def forward(self, x):
        x = x.unsqueeze(1).float()  # 增加通道维度：(batch, 10, 9) → (batch, 1, 10, 9)
        x = self.conv_layers(x)
        x = self.fc_layers(x)
        return x.squeeze(-1)


class MCTSNode:
    """蒙特卡洛树搜索节点"""
    def __init__(self, board_state, current_player, parent=None, move=None):
        self.board_state = board_state  # 棋盘状态（张量）
        self.current_player = current_player  # 当前玩家
        self.parent = parent  # 父节点
        self.move = move  # 到达该节点的落子
        self.children = []  # 子节点列表
        self.visit_count = 0  # 访问次数
        self.value = 0.0  # 累计价值
        self.untried_moves = self.get_all_valid_moves()  # 未尝试的合法移动

    def get_all_valid_moves(self):
        """获取当前棋盘的所有合法移动"""
        valid_moves = []
        # 创建临时棋盘用于计算有效移动
        temp_board = ChineseChessBoard()
        temp_board.board = self.board_state.clone()
        temp_board.current_player = self.current_player
        
        for y in range(BOARD_ROWS):
            for x in range(BOARD_COLS):
                if temp_board.board[y, x].item() * self.current_player > 0:
                    for (ny, nx) in temp_board.get_piece_moves(y, x):
                        valid_moves.append((y, x, ny, nx))
        return valid_moves

    def select_child(self, c=1.414):
        """UCB1算法选择子节点（平衡探索与利用）"""
        return max(self.children, key=lambda node: 
                  (node.value / node.visit_count) + 
                  c * np.sqrt(np.log(self.visit_count) / node.visit_count))

    def expand(self):
        """扩展子节点（生成新的落子节点）"""
        move = self.untried_moves.pop()
        from_y, from_x, to_y, to_x = move
        
        # 创建新棋盘状态
        new_board_state = self.board_state.clone()
        new_board_state[to_y, to_x] = new_board_state[from_y, from_x]
        new_board_state[from_y, from_x] = 0
        
        # 切换玩家
        new_player = -self.current_player
        
        child_node = MCTSNode(new_board_state, new_player, parent=self, move=move)
        self.children.append(child_node)
        return child_node

    def backpropagate(self, result):
        """回溯更新节点价值和访问次数"""
        self.visit_count += 1
        self.value += result
        if self.parent:
            self.parent.backpropagate(-result)  # 父节点视角相反


class ChessAI:
    """AI控制器：结合价值网络和MCTS实现决策"""
    def __init__(self, num_simulations=200):
        self.value_net = ChessValueNet()
        self.num_simulations = num_simulations  # MCTS模拟次数（影响AI强度）
        self.value_net.eval()  # 评估模式

    def predict_value(self, board_state):
        """用价值网络预测当前棋局胜率"""
        with torch.no_grad():
            return self.value_net(board_state).item()

    def mcts_search(self, board):
        """蒙特卡洛树搜索最优落子"""
        root = MCTSNode(board.board.clone(), board.current_player)
        for _ in range(self.num_simulations):
            node = root
            # 1. 选择：直到找到未扩展节点
            while not node.untried_moves and node.children:
                node = node.select_child()
            # 2. 扩展：生成子节点
            if node.untried_moves:
                node = node.expand()
            # 3. 评估：价值网络预测
            result = self.predict_value(node.board_state)
            # 4. 回溯：更新价值
            node.backpropagate(result)
        # 选择访问次数最多的子节点（最优落子）
        return max(root.children, key=lambda x: x.visit_count).move

    def get_ai_move(self, board):
        """获取AI的落子（对外接口）"""
        start_time = time.time()
        move = self.mcts_search(board)
        print(f"AI思考耗时：{time.time() - start_time:.2f}秒")
        return move


def main():
    """游戏主逻辑"""
    # 初始化游戏
    board = ChineseChessBoard()
    ai = ChessAI(num_simulations=300)  # 模拟次数越多AI越强（建议200-500）
    clock = pygame.time.Clock()
    game_over = False
    winner_text = ""

    while True:
        # 处理事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if not game_over and event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                # 玩家点击棋盘
                if board.handle_click(pygame.mouse.get_pos()):
                    # 检查AI是否被将死
                    if board.is_checkmate():
                        game_over = True
                        winner_text = "恭喜你！你赢了！"

        # AI回合
        if not game_over and board.current_player == -1:
            ai_move = ai.get_ai_move(board)
            from_y, from_x, to_y, to_x = ai_move
            board.move_piece(from_y, from_x, to_y, to_x)
            
            # 检查玩家是否被将死
            if board.is_checkmate():
                game_over = True
                winner_text = "很遗憾！AI赢了！"

        # 绘制棋盘
        board.draw_board()
        
        # 显示游戏结束信息
        if game_over:
            # 创建半透明覆盖层
            s = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
            s.fill((0, 0, 0, 128))
            screen.blit(s, (0, 0))
            
            # 显示获胜信息
            text = font.render(winner_text, True, WHITE)
            text_rect = text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2))
            screen.blit(text, text_rect)
            
            # 显示重新开始提示
            restart_text = font.render("按R键重新开始", True, WHITE)
            restart_rect = restart_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 40))
            screen.blit(restart_text, restart_rect)
            
            pygame.display.flip()
            
            # 检查重新开始
            keys = pygame.key.get_pressed()
            if keys[pygame.K_r]:
                board = ChineseChessBoard()
                game_over = False

        clock.tick(30)  # 限制帧率


if __name__ == "__main__":
    main()
